#!/usr/bin/python
# -*- coding: UTF-8 -*-
"""
The MIT License (MIT)
Copyright © 2022 Walkline Wang (https://walkline.wang)
Gitee: https://gitee.com/walkline/micropython-resource-manager

Append user resource files to specified partition of MicroPython firmware
"""
import os
import sys
from config import Config
from math import ceil

MKLITTLEFS_EXE = 'mklittlefs.exe'
LITTLEFS_PARTITION_TEMP = 'temp.bin'
BLOCK_SIZE = 1024 * 4


def show_help(self):
	print(
f'''
Usage:
    python3 {self} <firmware_file> <user_file_folder>
'''
	)

def __is_file_folder_exists(name, is_folder=False):
	if is_folder:
		return os.path.exists(name) and os.path.isdir(name)
	else:
		return os.path.exists(name) and os.path.isfile(name)

def __calculate_folder_blocks(folder):
	blocks = 0

	for dir in os.listdir(folder):
		fullpath = os.path.join(folder, dir)
		if os.path.isdir(fullpath):
			blocks += __calculate_folder_blocks(fullpath)
		else:
			blocks += ceil(os.stat(fullpath).st_size / BLOCK_SIZE)

	return blocks

def __generate_littlefs_partition(folder):
	blocks = __calculate_folder_blocks(folder) + 6 # 6 是块增补数量，没有合适的计算方式，只能估计

	try:
		print(blocks * BLOCK_SIZE)
		if os.system(f'{MKLITTLEFS_EXE} -p 256 -s {blocks * BLOCK_SIZE} -c {folder} {LITTLEFS_PARTITION_TEMP}') != 0:
			exit(1)
	except OSError as ose:
		print(ose)
		exit(1)

def main(argv):
	firmware, folder = argv

	if __is_file_folder_exists(firmware) and __is_file_folder_exists(folder, True):
		__generate_littlefs_partition(folder)

		file, ext = os.path.splitext(firmware)

		input_file = open(firmware, 'rb')
		output_file = open(f'{file}_output{ext}', 'wb')
		partition_file = open(LITTLEFS_PARTITION_TEMP, 'rb')

		# 复制固件内容到输出文件并移动指针到资源分区起始处
		output_file.write(input_file.read())
		output_file.seek(Config.Partition.RESOURCE_OFFSET)

		# 写入资源分区文件
		# _ = partition_file.seek(0x1000)
		output_file.write(partition_file.read())

		output_file.flush()
		output_file.close()
		input_file.close()
		partition_file.close()

		try:
			os.remove(LITTLEFS_PARTITION_TEMP)
		except:
			pass
	else:
		print(f'firmware_file or user_file_folder not exists')


if __name__ == '__main__':
	if not __is_file_folder_exists(MKLITTLEFS_EXE):
		print('missing mklittlefs.exe, download file by visiting https://github.com/earlephilhower/mklittlefs/releases')
		exit(0)

	argv = sys.argv

	if len(argv) in (1, 2) or len(argv) > 3:
		show_help(argv[0])
		exit(0)

	main(sys.argv[1:])
